import os
import re
from html import unescape
from typing import List
from urllib import parse

from Cryptodome.Cipher import AES
from lxml import html
from w3lib.html import remove_tags

from entity import Course, Chapter
from enums import CourseType
from sdk import decrypt


class Downloader:
    def __init__(self, client):
        self.client = client
        self.course_id = None

    def get_m3u8_url(self, chapter: Chapter) -> str:
        raise NotImplementedError()

    def get_m3u8_data(self, url: str) -> str:
        response = self.client.session.get(url)
        json_data = response.json()

        m3u8 = decrypt(json_data["data"]["info"], flag=0)

        for line in m3u8.splitlines():
            if line.startswith("https://"):
                response2 = self.client.session.get(line)
                json_data2 = response2.json()

                m3u8_data = decrypt(json_data2["data"]["info"], flag=0)
                return m3u8_data

        raise Exception(f"没有找到m3u8信息: {m3u8}")

    def get_m3u8_key(self, m3u8_data: str) -> bytearray:
        for line in m3u8_data.splitlines():
            if line.startswith("#EXT-X-KEY"):
                url = re.search(r"URI=\"(?P<url>.+)\"", line)["url"]

                response = self.client.session.get(url)
                json_data = response.json()

                key_dict: dict = decrypt(json_data["data"]["info"], flag=1)
                return bytearray(key_dict.values())

        raise Exception("没有找到key url")

    @staticmethod
    def get_m3u8_playlist(m3u8_data: str) -> List[str]:
        return [line for line in m3u8_data.splitlines() if line.startswith("https://")]

    def get_course(self) -> Course:
        raise NotImplementedError()

    def download_chapter(self, chapter: Chapter, key: bytearray, playlist: List[str], video_dir: str):
        with open(f"{video_dir}/{chapter.title.replace('/','')}.mp4", "wb") as f:
            aes = AES.new(key, AES.MODE_CBC)
            for url in playlist:
                response = self.client.session.get(url)
                data = aes.decrypt(response.content)
                f.write(data)

    def download(self, course_id):
        self.course_id = course_id

        course = self.get_course()
        print(course.title)

        video_dir = os.path.join("videos", course.title)
        if not os.path.exists(video_dir):
            os.makedirs(video_dir)

        for chapter in course.chapters:
            print("开始下载", chapter)
            m3u8_url = self.get_m3u8_url(chapter)
			
            m3u8_data = self.get_m3u8_data(m3u8_url)

            key = self.get_m3u8_key(m3u8_data)
            playlist = self.get_m3u8_playlist(m3u8_data)

            self.download_chapter(chapter, key, playlist, video_dir)	


class FreeCourseDownloader(Downloader):
    def get_course(self) -> Course:
        url = f"https://www.imooc.com/learn/{self.course_id}"
        response = self.client.session.get(url)

        html_doc: html.HtmlElement = html.fromstring(response.content)

        course_title = html_doc.xpath("//h2[@class='l']/text()")[0].strip()
        chapters = []

        for tag in html_doc.xpath("//ul[@class='video']/li[@data-media-id]"):
            a = tag.xpath("a[@class='J-media-item']")[0]

            url = parse.urljoin(url, a.xpath("@href")[0])
            title = unescape(remove_tags(html.tostring(a)).strip()).split("\n")[0].strip()
            chapters.append(Chapter(url=url, title=title))
		
        return Course(title=course_title, chapters=chapters)

    def get_m3u8_url(self, chapter: Chapter) -> str:
        response = self.client.session.get(chapter.url)
        mongo_id = re.search(rb"OP_CONFIG\.mongo_id=\"(\w+)\";", response.content)[1]

        url = "https://www.imooc.com/course/playlist/"
        params = {
            "t": "m3u8",
            "_id": mongo_id,
            "cdn": "aliyun1",
        }
        return url + "?" + parse.urlencode(params)


class PracticeCourseDownloader(Downloader):
    def get_course(self) -> Course:
        url = f"https://coding.imooc.com/learn/list/{self.course_id}.html"
        response = self.client.session.get(url)

        html_doc: html.HtmlElement = html.fromstring(response.content)

        course_title = html_doc.xpath("//h2[@class='course-title']/a/text()")[0].strip()

        chapters = []
        for tag in html_doc.xpath("//div[@class='list-item']/ul/li"):
            if not tag.xpath("i[contains(@class,'imv2-video')]"):
                continue

            url = parse.urljoin(url, tag.xpath("a/@href")[0])
            title = tag.xpath("a/span/text()")[0].strip()

            chapters.append(Chapter(url=url, title=title))

        return Course(title=course_title, chapters=chapters)

    def get_m3u8_url(self, chapter: Chapter) -> str:
        url = "https://coding.imooc.com/lesson/m3u8h5"
        mid = re.search(r"mid=(?P<mid>\d+)", chapter.url)["mid"]

        params = {
            "mid": mid,
            "cid": self.course_id,
            "ssl": 1,
            "cdn": "aliyun1",
        }
        return url + "?" + parse.urlencode(params)


class GenericDownloader:
    DOWNLOADER_CLASS_MAP = {
        CourseType.FREE: FreeCourseDownloader,
        CourseType.MINI: FreeCourseDownloader,
        CourseType.PRACTICE: PracticeCourseDownloader,
    }

    def __init__(self, client):
        self.client = client

    def download(self, course_id, course_type: CourseType):
        try:
            downloader = self.DOWNLOADER_CLASS_MAP[course_type](self.client)
        except KeyError as e:
            raise ValueError(f"未知的课程类型: {course_type}") from e
        else:
            return downloader.download(course_id)
